LLM推理入门指南②:深入解析KV缓存
在本系列文章《LLM推理入门指南①:文本生成的初始化与解码阶段》中,作者对Transformer解码器的文本生成算法进行了高层次概述,着重介绍了两个阶段:单步初始化阶段,即提示的处理阶段,和逐个生成补全词元的多步生成阶段。
本文进一步探讨了LLM推理的挑战 —— 第一大挑战是,注意力层(也称为自注意力层)与总序列长度(提示词元和生成补全词元)的计算成本呈二次方扩展的问题。幸运的是,生成步骤之间的许多计算是冗余的,从而能够缓存适当的结果并减少计算需求。
这种缓存方案也就是KV缓存,是LLM推理过程中的一种常用优化方式,使得(自)注意力机制的计算需求在总序列长度(提示 + 生成的完成部分)上线性扩展,而不是呈二次方扩展。更具体地说,KV缓存通过在生成过程中计算出的键和值张量存储(“缓存”)在GPU内存中,从而在每个生成步骤中节省了对过去词元键和值张量的重新计算。
KV缓存是一种折衷方案:以内存换取计算。本文后半部分将介绍KV缓存可以增长到多大,会带来哪些挑战,并介绍解决这些挑战的最常见应对策略。
(本文作者为AWS的GenAI解决方案架构师Pierre Lienhart。以下内容由OneFlow编译发布,转载请联系授权。原文:https://medium.com/@plienhar/llm-inference-series-3-kv-caching-unveiled-048152e461c8)
1
简要回顾Transformer注意力层
首先,我们来回顾一下vanilla Transformer(图1)中多头注意力(MHA)层的一些事实。
图1:Transformer解码器层的详细视图(上)和一个输入序列长度为3的双头(自)注意力层(下)
为简单起见,假设我们只处理长度为t的单个序列(即批处理大小为1):
在整个过程中,输入序列(提示)中的每个词元都由一个密集向量表示(图1的浅黄色部分)。 注意力层的输入是一个密集向量序列,每个输入词元都对应一个向量,这些向量由前一个解码器块生成。 对于每个输入向量,注意力层生成一个具有相同维度的单个密集向量(图1的浅蓝色部分)。
现在,考虑单个注意力头的情况:
首先,我们通过三种不同的投影(图1最左侧的浅灰色向量)为每个输入向量生成三个较低维度的密集向量:查询(query)、键(key)和值(value)。共有t个查询向量、t个键向量和t个值向量。 对于每个查询,我们生成一个输出向量,该向量等于值的线性组合,该线性组合的系数即为注意力分数。换句话说,对于每个查询,相应的输出向量是值的注意力加权平均值。对于给定查询,注意力分数是由该查询与每个键的点积得出的。这样我们就为序列中的每个词元生成了一个表示,其中包含来自其他词元的信息,意味着我们创建了每个词元的上下文表示。 然而,在自回归解码的上下文中,我们无法使用所有可能的值来构建给定查询的输出表示。实际上,在计算与特定词元相关的查询的输出时,我们不能使用序列中后续出现的词元的值向量。这一限制通过掩码(masking)来实现,将被禁止的值向量的注意力分数设置为零,即被禁止的词元的注意力分数设置为零。
最后,将每个注意力头的输出连接,并使用最后一个线性变换进行转换,以产生最终输出。
2
注意力计算的二次方扩展
我们看一下计算注意力分数所需的浮点运算数(FLOPs)。对于一个给定的注意力头,对于批处理大小为b,总长度为t的每个序列(包括提示和生成的补全部分),注意力分数矩阵是通过将形状为(t, d_head)的查询张量与形状为(d_head, t)的转置键张量相乘而创建的。
一个矩阵乘法所需的FLOPs是多少?将形状为(n, p)的矩阵与另一个形状为(n, m)的矩阵相乘大致需要 2.m.n.p 次操作。在这种情况下,单头单序列的注意力分数计算大约需要 2.d_head.t^2 FLOPs。因此,总体上注意力分数计算需要 2.b.n_layers.n_head.d_head.t^2=2.b.n_layers.d_model.t^2 FLOPs。现在,t 的二次方扩展变得清晰可见。
从具体数字来看,以Meta的Llama2–7B为例,n_layers=32 且 d_model=4096。
注意:使用值张量乘以掩码注意力分数矩阵所需的FLOPs与上述计算量相同。
那么涉及模型权重的矩阵乘法呢?通过类似的分析,可以证明其计算复杂度为 O(b.n_layers.d_model^2.t),即计算需求与总序列长度 t 呈线性增长。
下面我们通过一个例子来理解二次方扩展的严重性。如果要生成第1001个词元,模型必须执行比生成第101个词元多100倍的FLOPs。这种计算量的指数级增长显然很快会变得难以承受。幸运的是,由于掩码的使用,实际上在生成步骤之间可以节省大量计算。
3
掩码在生成阶段引发了冗余计算
现在来到问题的关键。由于掩码技术的使用,对于给定词元,输出表示仅使用先前词元的表示生成。由于先前的词元在迭代过程中保持不变,因此对于该特定词元的输出表示对于所有后续迭代也将是相同的,因此意味着存在冗余计算。
我们以上一篇博文中的词元序列为例(其特性为每个单词即为一个词元)。假设我们刚刚从输入序列“ What color is the sky? The sky ”中生成了“ is ”。在过去的迭代中,“sky ”是输入序列的最后一个词元,因此与该词元相关联的输出表示是使用序列中所有词元的表示生成的,即“ What”、“ color”、“ is”、“ the”、“ sky”、“?”、“ The ”和“sky ”的值向量。
下一次迭代的输入序列将是“ What color is the sky? The sky is ”,但由于使用了掩码,从“sky”的视角来看,它似乎仍然是“ What color is the sky? The sky ”的输入序列。因此,“sky”生成的输出表示将与上一次迭代相同。
现在看一个图示(见图2),使用图1中的图表。启动步骤应该处理长度为1的输入序列。冗余计算的部分在浅红色和浅紫色中突出显示。浅紫色元素对应冗余计算的键和值。
“is”的查询向量。 “What”、“ color”、“ is”、“ the”、“ sky”、“?”、“The ”、“sky ”和“ is ”的键向量,用于计算注意力分数。 “What”、“ color”、“ is”、“ the”、“ sky”、“?”、“The ”、“sky ”和“ is ”的值向量,用于计算输出。
计算“is”的查询、键和值。 从缓存中获取“ What”、“ color”、“ is”、“ the”、“ sky”、“?”、“The ”和“sky ”的键和值向量,并将它们与我们刚刚为“is”计算的键和值向量连接起来。 使用“is”的查询和所有键来计算注意力分数。 使用注意力分数和所有值来计算“is”的输出向量。
初始化阶段实际上不受KV缓存策略的影响,因为没有之前的步骤。
然而,对于解码阶段,情况现在看起来完全不同了。我们不再使用整个序列作为输入,而是只用最后生成的词元(和KV缓存)。
4
KV缓存实现了注意力的线性扩展
5
KV缓存的实际例子
6
KV缓存带来的新问题
KV缓存也可能会消耗大量的GPU内存。不幸的是,即使加载相对较小的LLM,GPU内存也是有限的。因此,当增加总序列长度(上下文窗口大小)或一次处理的序列数量(即吞吐量)时,KV缓存成为了主要的技术障碍,从而提高了成本效率。
与我们需要从内存中搬运的数据量相比,KV缓存大大减少了我们在单个生成步骤中执行的操作量:我们获取大型权重矩阵和不断增长的KV缓存,只是为了执行微不足道的矩阵到向量的操作。不幸的是,在现在的硬件上,我们最终花费的时间更多地用于数据加载,而不是实际的数值计算,这显然导致了GPU计算能力的低效利用。换句话说,我们的GPU利用率较低,因此成本效率低下。
7
KV缓存可以增长到多大?
减少模型权重内存占用(权重量化)
减少KV缓存内存占用(见下文)
通过将模型分片到多个GPU上池化来自多个设备的内存,以网络通信为代价(模型并行化),或者使用其他类型的存储,如CPU内存或磁盘(offloading)
8
减少批大小?
9
减少对总序列长度的依赖?
10
减少层数?
11
减少注意力头的数量如何?
12
注意力头的隐藏维度?
13
使用每个参数更少的字节?
14
高效内存管理的重要性
由于请求的总序列长度事先不确定,我们可能会预留连续的内存块,以容纳最大的序列长度。这种分配中的一部分几乎肯定永远不会被使用,但也无法被其他请求利用,最终被浪费(即内部内存碎片化)。 即使序列长度事先已知,由于内存逐渐被消耗,但内存块被保留到请求的整个生命周期,较短的请求仍然无法使用尚未使用的内存块。 如果我们采用生成多个序列的解码策略,如束搜索,多个候选序列实际上可能会部分共享它们的KV缓存。如果我们没有考虑到这一点,那么我们将不可避免地浪费内存,因为我们会存储本来可以共享的重复KV条目。
15
如果GPU内存不足,为什么不“简单地”使用多个GPU或者将负载转移到CPU内存甚至磁盘上?
16
总结
[1]: Llama 2: Open Foundation and Fine-Tuned Chat Models (Touvron et al., 2023)
[2]: OPT: Open Pre-trained Transformer Language Models (Zhang et al., 2022)
[3]: Release blog posts for: MPT-7B (May 2023) and MPT-30B (June 2023)
[4]: BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (BigScience, 2023)
[5]: Scaling Laws for Neural Language Models (Kaplan et al., 2020)
[6]: Mistral 7B (Jiang et al., 2023)
[7]: Efficient Streaming Language Models with Attention Sinks (Xiao et al., 2023) + GitHub repository
[8]: H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (Zhang et al., 2023) + GitHub repository
[9]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (Liu et al. 2023)
[10]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (Ge et al., 2023)
[11]: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
[12]: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
[13]: PaLM: Scaling Language Modeling with Pathways (Chowdhery et al., 2022)
[14]: The Falcon Series of Open Language Models (Almazrouei et al., 2023)
[15]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al., 2023) + GitHub repository
[16]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar et al., 2022) + GitHub repository
[17]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers et al., 2022) + GitHub repository
[18]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao et al., 2022) + GitHub repository
[19]: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (Sheng et al., 2023) + GitHub repository
[20] Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023) + GitHub repository
[21] vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention (Kwon et al. 2023)
[22] Efficiently Programming Large Language Models using SGLang (Zheng et al., 2023) + Blog post
[23]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang et al., 2018)
[24]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)
[25]: Efficiently Scaling Transformer Inference (Pope et al., 2022)
[26]: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin et al., 2024)
[2]: OPT: Open Pre-trained Transformer Language Models (Zhang et al., 2022)
[3]: Release blog posts for: MPT-7B (May 2023) and MPT-30B (June 2023)
[4]: BLOOM: A 176B-Parameter Open-Access Multilingual Language Model (BigScience, 2023)
[5]: Scaling Laws for Neural Language Models (Kaplan et al., 2020)
[6]: Mistral 7B (Jiang et al., 2023)
[7]: Efficient Streaming Language Models with Attention Sinks (Xiao et al., 2023) + GitHub repository
[8]: H_2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models (Zhang et al., 2023) + GitHub repository
[9]: Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time (Liu et al. 2023)
[10]: Model Tells You What to Discard: Adaptive KV Cache Compression for LLMs (Ge et al., 2023)
[11]: Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019)
[12]: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023)
[13]: PaLM: Scaling Language Modeling with Pathways (Chowdhery et al., 2022)
[14]: The Falcon Series of Open Language Models (Almazrouei et al., 2023)
[15]: AWQ: Activation-aware Weight Quantization for LLM Compression and Acceleration (Lin et al., 2023) + GitHub repository
[16]: GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers (Frantar et al., 2022) + GitHub repository
[17]: LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scale (Dettmers et al., 2022) + GitHub repository
[18]: SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models (Xiao et al., 2022) + GitHub repository
[19]: FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU (Sheng et al., 2023) + GitHub repository
[20] Efficient Memory Management for Large Language Model Serving with PagedAttention (Kwon et al., 2023) + GitHub repository
[21] vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention (Kwon et al. 2023)
[22] Efficiently Programming Large Language Models using SGLang (Zheng et al., 2023) + Blog post
[23]: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism (Huang et al., 2018)
[24]: Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Narayanan et al., 2021)
[25]: Efficiently Scaling Transformer Inference (Pope et al., 2022)
[26]: Infinite-LLM: Efficient LLM Service for Long Context with DistAttention and Distributed KVCache (Lin et al., 2024)
【语言大模型推理最高加速11倍】SiliconLLM是由硅基流动开发的高效、易用、可扩展的LLM推理加速引擎,旨在为用户提供开箱即用的推理加速能力,显著降低大模型部署成本,加速生成式AI产品落地。(技术合作、交流请添加微信:SiliconFlow01)
SiliconLLM的吞吐最高提升近4倍,时延最高降低近4倍
数据中心+PCIe:SiliconLLM的吞吐最高提升近5倍;消费卡场景:SiliconLLM的吞吐最高提升近3倍
System Prompt场景:SiliconLLM的吞吐最高提升11倍;MoE模型:推理 SiliconLLM的吞吐最高提升近10倍